# -*- coding: utf-8 -*-
'''
Created on 14.11.2019

@author: yu03
'''

import glob
import re
import datetime
import cv2
import matplotlib.pyplot as plt
from scipy import signal
from FFT_Interpolation import FFT_interpolation_2, line_cal, line_cal_fix
import numpy as np
import os
from scipy.optimize import curve_fit
from mpl_toolkits.mplot3d import Axes3D
from PyUeye_Unified.Hor_16line import folder_path, np_names, hor_index, hor_lines
import sys

calculation_mode = 'diff' ### DC removed
# calculation_mode = 'DC' ### DC Not removed


def fit_func(x, a, b, c):
    return a*(x-b)**2 + c

''' 
    General Parameters
'''
fs_cam = 1000
Lamda = 633e-9
pix_size = 5.3e-6
V_x, V_y, V_z = 0, 0, 0


# for video in video_name:
#     i += 1
#     file_name = re.findall(r"Compare\\(.+).avi", video)[0]
#     np_result_name = folder_path + '\\' + file_name[0] + '_lines' + '.npy'
#     print(video)
#     print(i, file_name)

#     print(txt_name)


for np_file in np_names:
    ''' 
        File name define
    '''
    file_name = re.findall(r"Simulated_File\\(.+).npy", np_file)[0]
    if calculation_mode == 'DC':
        np_result_name = folder_path + '\\results\\' + file_name + '_lines' + '.npy'
    elif calculation_mode == 'diff':
        np_result_name = folder_path + '\\results\\diff\\' + file_name + '_lines_diff' + '.npy'
    else:
        sys.exit('FFT Mode Error:\n calculation_mode = %s'%calculation_mode)
    print('File:', file_name, '; ', datetime.datetime.now().time())
#     print(np_result_name)
    ''' 
        Reading .npy file
    '''
    f = open(np_file, 'rb')
    f_size = os.fstat(f.fileno()).st_size
    frames = []
    time_sequence = []
    while f.tell() < f_size:
    #         print(f.tell())
        frame = np.load(f, allow_pickle=True)
        frames.append(frame[0])
        time_sequence.append(frame[1])
    
    # x_axis = np.linspace(0, len(time_sequence)-1, len(time_sequence))
    
    frames = np.array(frames)
    print(frames.shape)
    img_set = frames
    frame_total = len(frames)
    line_num = len(frames[0])
    ''' 
        Finding Phase Spectrum Component
    '''
    hor_m_k_num_set = []
    ver_m_k_num_set = []
    for img in img_set[::5]:
        hor_center = np.array(img[0])
        hor_center = np.array([i[0]+256*i[1] for i in hor_center])
        if calculation_mode == 'diff':
            hor_center = np.diff(hor_center.astype('int'))
            hor_center = np.concatenate(([0], hor_center))
        elif calculation_mode == 'DC':
            pass
        else:
            sys.exit('FFT Mode Error:\n calculation_mode = %s'%calculation_mode)
        
        hor_f_fit, hor_phase_estim, hor_m_k_num = line_cal(hor_center)
        
        hor_m_k_num_set.append(hor_m_k_num)

    m_k_num_hor = int(np.average(hor_m_k_num_set))
    print(frame_total, "Frames;  ", 'Chosen Components:', m_k_num_hor, '  ', datetime.datetime.now().time() )
    
    ''' 
        Processing
    '''    
    hor_f_set = []
    hor_phi_set = []
    frame_num = 0
    for img in img_set:
#     for img in img_set[0:500]:
        frame_num += 1
        for i in range(line_num):
            hor_center = np.array(img[hor_lines[i]])
            hor_center = np.array([i[0]+256*i[1] for i in hor_center])
            
            if calculation_mode == 'diff':
                hor_center = np.diff(hor_center.astype('int'))
                hor_center = np.concatenate(([0], hor_center))
            elif calculation_mode == 'DC':
                pass
            else:
                sys.exit('FFT Mode Error:\n calculation_mode = %s'%calculation_mode)
            
            hor_f_fit, hor_phase_estim, hor_m_k_num = line_cal_fix(hor_center, m_k_num_hor)
               
            hor_f_set.append(hor_f_fit)
            hor_phi_set.append(hor_phase_estim)
#             ver_f_set.append(ver_f_fit)
#             ver_phi_set.append(ver_phase_estim)
        if frame_num%10 == 0:
            print('Frame: ', frame_num, '   ',datetime.datetime.now().time())   
    hor_f_set = np.array(hor_f_set)
    hor_f_set = hor_f_set.reshape(line_num, frame_num, order='F')
  
    hor_phi_set = np.array(hor_phi_set)
    hor_phi_set = hor_phi_set.reshape(line_num, frame_num, order='F')
  
    hor_angle_set = (V_x-Lamda*hor_f_set/2)*1e6 ### urad
    hor_length_set = np.unwrap(hor_phi_set)/4/np.pi*Lamda*1e9 ### nm
       
    print('line, frame:',np.shape(hor_angle_set))
    print('')
    f = open(np_result_name,'ab')
    np.save(f, hor_angle_set, allow_pickle=True)
    np.save(f, hor_length_set, allow_pickle=True)

''' 
    Plotting
'''   
y = [np.arange(frame_num)]*line_num
y = np.array(y)
   
fig = plt.figure('hor_tilt')
x = [hor_lines]*frame_num
x = np.array(x).T
ax = fig.gca(projection='3d')
plt.gca().patch.set_facecolor('white')
for i in range(line_num):
    ax.plot(x[i,:],y[i,:],hor_angle_set[i,:])

      
fig = plt.figure('hor_length')
x = [hor_lines]*frame_num
x = np.array(x).T
ax = fig.gca(projection='3d')
plt.gca().patch.set_facecolor('white')
for i in range(line_num):
    ax.plot(x[i,:],y[i,:],hor_length_set[i,:])
  
plt.show()